import matplotlib.pyplot as plt
from collections import defaultdict

# ======================
# 1. 你的 Z 偏移检测函数（已更新逻辑）
# ======================

def detect_z_shift(data, x_limit, y_limit):
    sorted_data = sorted(data, key=lambda item: item[0])
    n = len(sorted_data)

    segment_size = n // 4
    segments = [
        sorted_data[i * segment_size : (i + 1) * segment_size]
        for i in range(4)
    ]

    averages = []
    for seg in segments:
        if len(seg) == 0:
            avg = 0.0
        else:
            avg = sum(y for x, y in seg) / len(seg)
        averages.append(avg)
    
    a, b, c, d = averages

    condition3 = a - c - y_limit     # 第1段比第3段高出的值 减去限制
    condition4 = a - d - y_limit
    condition5 = b - c - y_limit
    condition6 = b - d - y_limit

    # 判断所有差值是否都为正（前高后低 -> 下降型 Z 偏移）
    if condition3 > 0 and condition4 > 0 and condition5 > 0 and condition6 > 0:
        kk = "down"  # 前高后低，显著下降
    # 判断所有差值是否都为负（前低后高 -> 上升型 Z 偏移）
    elif condition3 < 0 and condition4 < 0 and condition5 < 0 and condition6 < 0:
        kk = "up"    # 前低后高，显著上升
    else:
        kk = ""      # 无明显一致偏移

    return {
        "a": a, "b": b, "c": c, "d": d,
        "x_limit": x_limit,
        "y_limit": y_limit,
        'kk': kk,                   # 分类：'down', 'up', ''
        "is_z_shift_detected": bool(kk),  # 是否检测到明确的偏移类型
        "segments": segments
    }

# ======================
# 2. 你的示例数据与参数
# ======================

# 前20个点：x 从 0 到 19，y值较低
part1 = [(i, 10 + (i % 4)) for i in range(20)]

# 后20个点：x 从 20 到 39，y值较高
part2 = [(i, 13 + (i % 4)) for i in range(20, 40)]

# 合并数据
data = part1 + part2

x_limit = 1.0    # 小区间均值差异限制
y_limit = 2.0    # 大区间差值限制（用于判断是否显著高于/低于）

# ======================
# 3. 调用检测函数
# ======================

result = detect_z_shift(data, x_limit, y_limit)

# 打印检测结果
print("各段均值 a, b, c, d:", result["a"], result["b"], result["c"], result["d"])
print("x 限制（小区间）:", result["x_limit"])
print("y 限制（大区间）:", result["y_limit"])
print("检测结果分类 kk:", result["kk"])  # 'down'、'up' 或 ''

# 注意：你已不再返回 conditions 字典，所以这行需要删除或注释掉 ⬇️
# print("6个条件判断结果:")  # ❌ 原代码中用了 result["conditions"]，但你已移除
# for name, val in result["conditions"].items():  # ❌ 会报错 KeyError
#     print(f"  {name}: {'✓' if val else '✗'}")

print("👉 是否检测到 Z 偏移？", "是（类型：" + result["kk"] + "）" if result["is_z_shift_detected"] else "否")

# ======================
# 4. 绘图：数据 + 分段 + 均值线
# ======================

plt.figure(figsize=(12, 6))

segments = result["segments"]
colors = ['blue', 'green', 'orange', 'red']
labels = ['第1段 (0~9)', '第2段 (10~19)', '第3段 (20~29)', '第4段 (30~39)']

# 绘制每个分段
for i in range(4):
    seg = segments[i]
    x_vals = [p[0] for p in seg]
    y_vals = [p[1] for p in seg]
    plt.plot(x_vals, y_vals, marker='o', linestyle='-', color=colors[i], label=labels[i], markersize=4)

# 计算每段的 x 范围
segment_ranges = [
    (0, 9),      # 第1段
    (10, 19),    # 第2段
    (20, 29),    # 第3段
    (30, 39)     # 第4段
]

means = [result["a"], result["b"], result["c"], result["d"]]
y_positions = means

# 绘制均值水平线 + 数值标注
for i in range(4):
    y_mean = means[i]
    x_start, x_end = segment_ranges[i]
    x_pos = (x_start + x_end) / 2

    plt.axhline(y=y_mean, color=colors[i], linestyle='--', alpha=0.6)
    plt.text(x_pos, y_mean + 0.1, f'μ={y_mean:.2f}', color=colors[i], 
             fontsize=9, ha='center', weight='bold')
    plt.scatter(x_pos, y_mean, color=colors[i], s=50, zorder=5, edgecolor='black')

# 图表装饰
plt.title(f"数据分段可视化 & Z偏移检测 (结果类型: {result['kk']})", fontsize=14)  # 显示检测类型
plt.xlabel("X 轴（样本序号 / 位置）", fontsize=12)
plt.ylabel("Y 轴（测量值 / 可能的 Z 值）", fontsize=12)
plt.legend(loc='upper right')
plt.grid(True, alpha=0.3)
plt.ylim(8, 15)  # 根据数据范围调整
plt.tight_layout()
plt.show()
